AWS Amplify Gen 2 で Cognito のトークン生成前 Lambda トリガー（トリガーイベント V2）を使ってみた

いわさ
2024.08.26
いわさです。
AWS Amplify では標準の認証に Cognito ユーザープールを使用しています。

Cognito ユーザープールでは Lambda トリガーの機能を使って認証フローをカスタマイズすることが出来ます。
https://dev.classmethod.jp/articles/cognito-presignin-custom/
標準というか Amplify Auth L3 コントラクトのトリガープロパティでも様々なことが行えるのですが、さらに高度なセキュリティ機能を有効化することで追加のカスタマイズを行うことが出来ます。

具体的にはトークン生成前トリガーでアクセストークンのカスタマイズなどを行うことが出来ます。
https://dev.classmethod.jp/articles/cognito-user-pools-customize-access-tokens/
今回 Amplify Gen 2 でトークン生成前イベントの Lambda トリガー（V2）を有効化する機会があったので方法を紹介します。

本日時点の Amplify Backend の機能としてはまだサポートされていないので Amplify が生成するリソースに色々と追加の設定をしてやる感じになります。
 高度なセキュリティを有効化するまずは Cognito ユーザープールの高度なセキュリティ機能を有効化するのですが、以下を応用して Cognito ユーザープールリソースに直接アクセスします。
https://dev.classmethod.jp/articles/enable-delete-protection-amplify-auth-data/
ユーザープールの API や CloudFormation の場合はUserPoolAddOnsプロパティのAdvancedSecurityModeを設定してやれば良いので、その方法をそのまま使います。
amplify/backend.ts
import { defineBackend } from '@aws-amplify/backend';
import { auth } from './auth/resource';
import { data } from './data/resource';

const backend = defineBackend({
  auth,
  data,
});

const { cfnUserPool } = backend.auth.resources.cfnResources
cfnUserPool.userPoolAddOns = {
  advancedSecurityMode: "AUDIT"
};
高度なセキュリティ機能の有効化自体はこれだけで良いので、ここは非常に簡単です。
 トークン生成前 Lambda トリガーを設定する続いて Lambda トリガーを設定します。

Amplify で Lambda トリガーを設定する方法は次の記事を応用します。

Amplify Auth には Lambda トリガーを設定する仕組みが用意されているのでそいつをまずは使ってみます。
https://dev.classmethod.jp/articles/amplify-signup-domain-filter/
amplify/auth/pre-token-generation/handler.ts
import type { PreTokenGenerationTriggerHandler } from 'aws-lambda';
import { env } from '$amplify/env/pre-sign-up';

export const handler: PreTokenGenerationTriggerHandler = async (event) => {
  event.response = {
    claimsOverrideDetails: {
      claimsToAddOrOverride: {
        my_first_attribute: "true",
        my_second_attribute: "111",
      },
      claimsToSuppress: ["email"],
    },
  }

  return event;
};
amplify/auth/resource.ts
import { defineAuth } from '@aws-amplify/backend';
import { preSignUp } from './pre-sign-up/resource';
import { preTokenGeneration } from './pre-token-generation/resource';

/**
 * Define and configure your auth resource
 * @see https://docs.amplify.aws/gen2/build-a-backend/auth
 */
export const auth = defineAuth({
  loginWith: {
    email: true,
  },
  triggers: {
    preSignUp,
    preTokenGeneration,
  },
  userAttributes: {
    "custom:tenant_id": {
      dataType: "String",
      mutable: true,
    },
  },
});
authのtriggers.preTokenGenerationに設定してやります。

試してみたところ、高度なセキュリティ機能は有効化されていますが、V1_0 の Lambda トリガーとして設定されてしまっていました。ふむ...
 ユーザープールの V2 用プロパティに直接指定することが出来るCognito ユーザープールのトークン生成前トリガーですが、V1 と V2 で API 上の設定箇所が異なっています。

V1 はPreTokenGenerationに関数 ARN を指定し、V2 はPreTokenGenerationConfig内に関数 ARN とバージョン（V2_0）を指定します。なんか重複している気がするのでこのあたりは今後変わるかもと思いました。
https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html#user-pool-lambda-pre-token-generation-accesstoken
リファレンスから、PreTokenGenerationConfig.LambdaVersionをV2_0に設定したいところ。

Amplify のソースコードを眺めてみたのですが、Auth コントラクトだと対応していない雰囲気を感じました。

ということで今回は先ほどの高度なセキュリティ機能を有効化した時のように Amplify が作成するリソースの構成を追加してやります。
amplify/backend.ts
import { defineBackend } from '@aws-amplify/backend';
import { auth } from './auth/resource';
import { data } from './data/resource';

const backend = defineBackend({
  auth,
  data,
});

const { cfnUserPool } = backend.auth.resources.cfnResources
cfnUserPool.userPoolAddOns = {
  advancedSecurityMode: "AUDIT"
};
cfnUserPool.lambdaConfig = {
  preTokenGenerationConfig: {
    lambdaVersion: 'V2_0',
  }
};
まず、上記のようにpreTokenGenerationConfig.lambdaVersionだけ指定したところ次のようなエラーとなりました。
The CloudFormation deployment has failed.
Caused By: Deployment failed: Error: The stack named amplify-amplifyvitereacttemplate-iwasatakahito-sandbox-c16ed1c1a8 failed to deploy: UPDATE_ROLLBACK_COMPLETE: Resource handler returned message: "1 validation error detected: Value null at 'lambdaConfig.preTokenGenerationConfig.lambdaArn' failed to satisfy constraint: Member must not be null (Service: CognitoIdentityProvider, Status Code: 400, Request ID: 1d95e590-e32d-4e26-b765-f20d2767aa23)" (RequestToken: 804342e9-a460-0381-2b0b-e591ecbcef5a, HandlerErrorCode: InvalidRequest), Embedded stack arn:aws:cloudformation:ap-northeast-1:550669467088:stack/amplify-amplifyvitereacttemplate-iwasatakahito-sandbox-c16ed1c1a8-auth179371D7-2DPWAPKPNSTR/865dfa10-60ce-11ef-bcf2-067b1c30b61b was not successfully updated. Currently in UPDATE_ROLLBACK_IN_PROGRESS with reason: The following resource(s) failed to update: [amplifyAuthUserPool4BA7F805]. 
このパラメータはlambdaConfig.preTokenGenerationConfig.lambdaArnも一緒に設定してやる必要があるようです。
 V2 設定出来たが他のトリガーが消えたなので、今回は関数を個別に定義してバックエンドリソースに追加してやる感じにしました。

Auth のトリガーとしてではなくて普通の関数として設定した感じです。

その上でユーザープールに設定してやりました。
amplify/backend.ts
import { defineBackend } from '@aws-amplify/backend';
import { auth } from './auth/resource';
import { data } from './data/resource';
import { preTokenGenerationV2 } from './auth/pre-token-generation-v2/resource';

const backend = defineBackend({
  auth,
  data,
  preTokenGenerationV2,
});

const { cfnUserPool } = backend.auth.resources.cfnResources
cfnUserPool.userPoolAddOns = {
  advancedSecurityMode: "AUDIT"
};
cfnUserPool.lambdaConfig = {
  preTokenGenerationConfig: {
    lambdaVersion: 'V2_0',
    lambdaArn: backend.preTokenGenerationV2.resources.lambda.functionArn,
  }
};
なお、V2 は関数ハンドラーの構造が異なるのでPreTokenGenerationV2TriggerHandlerを使ってやる必要があります。
amplify/auth/pre-token-generation-v2/handler.ts
import type { PreTokenGenerationV2TriggerHandler } from 'aws-lambda';

export const handler: PreTokenGenerationV2TriggerHandler = async (event) => {
  event.response = {
    claimsAndScopeOverrideDetails: {
      idTokenGeneration: {
        claimsToAddOrOverride: {
          hoge1: "hoge1-id",
          hoge2: "hoge2-id",
        },
        claimsToSuppress: ["email"]
      },
      accessTokenGeneration: {
        claimsToAddOrOverride: {
          hoge1: "hoge1-access",
          hoge2: "hoge2-access",
        }
      }
    }
  }

  return event;
};
デプロイすると次のように V2 トリガーが設定出来ました。やった！
ただ、Auth の機能で設定したサインアップ前トリガーが消えてしまいました。
 必要な設定だけ addPropertyOverride で更新するLambdaConfigをすべて置き換えてしまっているので、標準で設定されるPreSignUpが消えてしまったようです。

そこで、次の方法を使って必要な部分だけ設定を上書きしてみることに。
https://dev.classmethod.jp/articles/amplify-gen2-resource-override/
amplify/auth/resource.ts
:

const { cfnUserPool } = backend.auth.resources.cfnResources
cfnUserPool.userPoolAddOns = {
  advancedSecurityMode: "AUDIT"
};
cfnUserPool.addPropertyOverride("LambdaConfig.PreTokenGenerationConfig",{
  LambdaVersion: 'V2_0',
  LambdaArn: backend.preTokenGenerationV2.resources.lambda.functionArn,
});
デプロイしてみるとサインアップ前 Lambda トリガーもちゃんと設定されていました。

良さそうです。
 Cognito へのリソースベースポリシーが設定されていないめでたしと思って確認してみたところ、Lambda トリガーの実行に失敗しました。

AccessDinedException...
Lambda 関数の権限関係の設定を見てみたところ、Auth コントラクトで設定した場合はリソースベースポリシーが設定されるのですが、関連付けせずに個別に作成した関数はリソースベースポリシーが設定されていませんでした。そりゃそうか。
個別に関数にリソースベースポリシーを設定してやるのが正しそうな気がするのですが、今回は次のように Auth 上に個別で作成した関数も設定してみることにしました。
amplify/auth/resource.ts
import { defineAuth } from '@aws-amplify/backend';
import { preSignUp } from './pre-sign-up/resource';
// import { preTokenGeneration } from './pre-token-generation/resource';
import { preTokenGenerationV2 } from './pre-token-generation-v2/resource';

/**
 * Define and configure your auth resource
 * @see https://docs.amplify.aws/gen2/build-a-backend/auth
 */
export const auth = defineAuth({
  loginWith: {
    email: true,
  },
  triggers: {
    preSignUp,
    preTokenGeneration: preTokenGenerationV2,
  },
  userAttributes: {
    "custom:tenant_id": {
      dataType: "String",
      mutable: true,
    },
  },
});
デプロイしてみるとリソースペースポリシーが設定されていました。ただ、この設定方法はちょっと違う気もしているので見直したいところです。
 実行してアクセストークンを確認してみるバックエンドリソースが一通りデプロイ出来たので、サインインしてクライアントが使いまわしているアクセストークンの中身を jwt.io などで覗いてみましょう。
おっ、Lambda トリガーで設定したカスタムクレームが設定されていますね。
 さいごに本日は AWS Amplify Gen 2 で Cognito のトークン生成前 Lambda トリガー（トリガーイベント V2）を使ってみました。
まだ Amplify Auth コントラクトでは色々とサポートされていないようでしたが、L1・L2 あたりで追加設定することで対応出来ました。

トークンカスタマイズは結構需要がある気がするので、Auth でサポートされてほしいですね。